import pandas as pd
import jieba
from sklearn.feature_extraction.text import CountVectorizer
from snownlp import SnowNLP
from sklearn.naive_bayes import MultinomialNB
from sklearn.tree import DecisionTreeClassifier
from sklearn.model_selection import train_test_split

data = pd.read_csv('data/data1.csv')
data['star'].unique()


def make_label(star):
    if star > 3:
        return 1
    else:
        return 0

def chinese_word_cut(mytext):
    return " ".join(jieba.cut(mytext))


data['sentiment'] = data.star.apply(make_label)
data['cut_comment'] = data.comment.apply(chinese_word_cut)

X = data['cut_comment']
y = data.sentiment
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=22)


def get_custom_stopwords(stop_words_file):
    with open(stop_words_file) as f:
        stopwords = f.read()
    stopwords_list = stopwords.split('\n')
    custom_stopwords_list = [i for i in stopwords_list]
    return custom_stopwords_list


stop_words_file = 'data/哈工大停用词表.txt'
stopwords = get_custom_stopwords(stop_words_file)

vect = CountVectorizer(max_df=0.8,
                       min_df=3,
                       token_pattern=u'(?u)\\b[^\\d\\W]\\w+\\b',
                       stop_words=frozenset(stopwords))

test = pd.DataFrame(vect.fit_transform(X_train).toarray(), columns=vect.get_feature_names())

model = MultinomialNB()
# model = DecisionTreeClassifier()
X_train_vect = vect.fit_transform(X_train)
model.fit(X_train_vect, y_train)
train_score = model.score(X_train_vect, y_train)
print(train_score)


X_test_vect = vect.transform(X_test)
print(model.score(X_test_vect, y_test))

X_vec = vect.transform(X)
nb_result = model.predict(X_vec)
data['nb_result'] = nb_result


